//==========================================================================
// Copyright (c) 2000-2008,  Elastos, Inc.  All Rights Reserved.
//==========================================================================

#include <winhack.h>
#include <elapi.h>
#include <driver.h>
#include <dbgstub.h>

//
// CDebuggerStub
//
CDebuggerStub::CDebuggerStub() :
    m_dwTargetProcessId(0),
    m_hTargetProcess(0)
{
}

ECode CDebuggerStub::Read(
    /* [in] */ Int64 u64Offset,
    /* [in] */ Int32 bytesToRead,
    /* [out] */ MemoryBuf * pBuffer,
    /* [out] */ IEvent * * ppCompletionEvent)
{
    return E_NOT_IMPLEMENTED;
}

ECode CDebuggerStub::Write(
    /* [in] */ Int64 u64Offset,
    /* [in] */ const MemoryBuf & buffer,
    /* [out] */ Int32 * pBytesWritten,
    /* [out] */ IEvent * * ppCompletionEvent)
{
    return E_NOT_IMPLEMENTED;
}

ECode CDebuggerStub::Control(
        /* [in] */ Handle32 nControlCode,
        /* [in] */ const MemoryBuf & inBuffer,
        /* [out] */ MemoryBuf * pOutBuffer,
        /* [out] */ IEvent * * ppCompletionEvent)
{
    BOOL ret;
    ECode ec = NOERROR;
    UInt32 * pInArg = (UInt32 *)inBuffer.GetPayload();
    UInt32 nInArgSize = inBuffer.GetUsed();  // the number of bytes

    switch (nControlCode) {
        case DBGSTUB_CTRL_CODE_ATTACH_PROCESS: {
            if (inBuffer.IsNull() || nInArgSize < sizeof(Int32)) {
                return E_INVALID_ARGUMENT;
            }

            // Open the target process object by the process id
            //
            DWORD dwProcessId = *pInArg;
            m_hTargetProcess = OpenProcess(
                    PROCESS_ALL_ACCESS, FALSE, dwProcessId);
            if (NULL == m_hTargetProcess) {
                return GetLastErrorCode();
            }

            // Tto attach the target process and debug it
            //
            ret = DebugActiveProcess(dwProcessId);
            if (!ret) {
                ec = GetLastErrorCode();
                CloseHandle(m_hTargetProcess);
                m_hTargetProcess = NULL;
                return ec;
            }

            m_dwTargetProcessId = dwProcessId;
            return NOERROR;
        }
        break;

        case DBGSTUB_CTRL_CODE_DETACH_PROCESS: {

            // Stop the debugger from debugging the target process
            //
            //ret = DebugActiveProcessStop(m_dwTargetProcessId);

            ret = TerminateProcess(m_hTargetProcess, 0);
            if (!ret) {
                return GetLastErrorCode();
            }
            CloseHandle(m_hTargetProcess);
            m_hTargetProcess = NULL;
            m_dwTargetProcessId = 0;
            return NOERROR;
        }
        break;

        case DBGSTUB_CTRL_CODE_WAIT_FOR_DEBUGEVENT: {
            if (inBuffer.IsNull() || nInArgSize < (Int32)sizeof(DWORD)
                    || pOutBuffer->IsNull()
                    || pOutBuffer->GetCapacity() < (Int32)sizeof(DEBUG_EVENT)) {
                return E_INVALID_ARGUMENT;
            }

            // Wait for a debugging event to occur in the target process
            //
            ret = WaitForDebugEvent(
                    (LPDEBUG_EVENT)pOutBuffer->GetPayload(),
                    *(DWORD *)pInArg);          // milliseconds to wait
            if (!ret) {
                return GetLastErrorCode();
            }

            pOutBuffer->SetUsed(sizeof(DEBUG_EVENT));
            return NOERROR;
        }
        break;

        case DBGSTUB_CTRL_CODE_CONTINUE: {
            if (inBuffer.IsNull() || nInArgSize < (Int32)(sizeof(DWORD) * 2)) {
                return E_INVALID_ARGUMENT;
            }

            // Continue the thread that previously reported a debugging event
            //
            ret = ContinueDebugEvent(
                        m_dwTargetProcessId,
                        pInArg[0],      // Thread ID
                        pInArg[1]);     // continuation status
            if (!ret) {
                return GetLastErrorCode();
            }
            return NOERROR;
        }
        break;

        case DBGSTUB_CTRL_CODE_SET_THREAD_CONTEXT: {
            if (inBuffer.IsNull()
                    || nInArgSize < (Int32)(sizeof(Context) + sizeof(DWORD))) {
                return E_INVALID_ARGUMENT;
            }

            // Set the context for the target thread
            //
            HANDLE hThread = OpenThread(
                    THREAD_ALL_ACCESS, FALSE, *(DWORD *)pInArg);
            if (NULL == hThread) {
                return GetLastErrorCode();
            }

            ret = SetThreadContext(hThread, (Context *)(pInArg + 1));
            if (!ret) {
                ec = GetLastErrorCode();
            }
            else {
                ec = NOERROR;
            }

            CloseHandle(hThread);
            return ec;
        }
        break;

        case DBGSTUB_CTRL_CODE_GET_THREAD_CONTEXT: {

            if (inBuffer.IsNull()
                    || pOutBuffer->IsNull()
                    || pOutBuffer->GetCapacity()  < (Int32)sizeof(Context)) {
                return E_INVALID_ARGUMENT;
            }

            // Retrieve the context of the target thread
            //
            HANDLE hThread = OpenThread(
                    THREAD_ALL_ACCESS, FALSE, *(DWORD *)pInArg);
            if (NULL == hThread) {
                return GetLastErrorCode();
            }

            Context * pContext = (Context *)pOutBuffer->GetPayload();
            pContext->ContextFlags = CONTEXT_DEBUG_REGISTERS
                    | CONTEXT_SEGMENTS | CONTEXT_INTEGER | CONTEXT_CONTROL;
            ret = GetThreadContext(hThread, pContext);
            if (!ret) {
                ec = GetLastErrorCode();
            }
            else {
                ec = NOERROR;
            }

            CloseHandle(hThread);

            if (SUCCEEDED(ec)) {
                pOutBuffer->SetUsed(sizeof(Context));
            }
            return ec;
        }
        break;

        case DBGSTUB_CTRL_CODE_READ_PROCESS_MEMORY: {
            if (inBuffer.IsNull()
                    || nInArgSize < (Int32)sizeof(Int32) * 2) {
                return E_INVALID_ARGUMENT;
            }

            PVoid base = (PVoid)pInArg[0];
            size_t size = pInArg[1];

            if (pOutBuffer->IsNull() || pOutBuffer->GetCapacity() < (Int32)size) {
                return E_INVALID_ARGUMENT;
            }

            // Read data from an area of memory in a the target process
            //
            PVoid pBuf = (PVoid)pOutBuffer->GetPayload();
            size_t uRead;

            ret = ReadProcessMemory(
                    m_hTargetProcess, base, pBuf, size, &uRead);
            if (!ret) {
                return GetLastErrorCode();
            }
            pOutBuffer->SetUsed(uRead);
            return NOERROR;
        }
        break;

        case DBGSTUB_CTRL_CODE_WRITE_PROCESS_MEMORY: {
            if (nInArgSize < (Int32)(2 * sizeof(Int32))
                    || pOutBuffer->IsNull()
                    || pOutBuffer->GetCapacity() < (Int32)sizeof(Int32)) {
                return E_INVALID_ARGUMENT;
            }

            PVoid pAddress = (PVoid)pInArg[0];
            size_t size = pInArg[1];
            PVoid pBuf = pInArg + 2;

            if (nInArgSize < size + 8) {
                return E_INVALID_ARGUMENT;
            }

            ec = this->WriteProcessMemory(m_hTargetProcess, pAddress, pBuf,
                    size, (size_t *)(PVoid)pOutBuffer->GetPayload());
            if (SUCCEEDED(ec)) {
                pOutBuffer->SetUsed(sizeof(Int32));
            }
            return ec;
        }
        break;

        case DBGSTUB_CTRL_CODE_SET_BREAKPIONT: {
            if (inBuffer.IsNull() || nInArgSize < (Int32)sizeof(Int32)
                    || pOutBuffer->IsNull()
                    || pOutBuffer->GetCapacity() < DBG_INS_LENGTH) {
                return E_INVALID_ARGUMENT;
            }

            HANDLE hThread = OpenThread(
                    THREAD_ALL_ACCESS, FALSE, *(DWORD *)pInArg);
            if (NULL == hThread) {
                return GetLastErrorCode();
            }

            Context ctx;
            ctx.ContextFlags = CONTEXT_ALL;
            ret = GetThreadContext(hThread, &ctx);
            if (!ret) {
                ec = GetLastErrorCode();
            }
            else {
                ec = NOERROR;
            }
            CloseHandle(hThread);
            if (FAILED(ec)) {
                return ec;
            }

            PVoid pAddress = (PVoid)ctx.Eip;

            ret = ReadProcessMemory(m_hTargetProcess, pAddress,
                    (PVoid)pOutBuffer->GetPayload(), DBG_INS_LENGTH, NULL);
            if (!ret) {
                return GetLastErrorCode();
            }

            pOutBuffer->SetUsed(DBG_INS_LENGTH);

            DWORD bps = 0xcc;
            return this->WriteProcessMemory(
                    m_hTargetProcess, pAddress, &bps, DBG_INS_LENGTH, NULL);
        }

        case DBGSTUB_CTRL_CODE_GET_THREAD_STACK_BASE: {
            if (inBuffer.IsNull() || nInArgSize < (Int32)sizeof(Int32)
                    || pOutBuffer->IsNull()
                    || pOutBuffer->GetCapacity() < (Int32)sizeof(PVoid)) {
                return E_INVALID_ARGUMENT;
            }

            ec = this->GetThreadStackBase(
                    *(DWORD *)pInArg,                   // process Id
                    *(DWORD *)(pInArg + 1),             // thread Id
                    (PVoid *)pOutBuffer->GetPayload());   // PVoid * ppStackBase
            if (SUCCEEDED(ec)) {
                pOutBuffer->SetUsed((Int32)sizeof(PVoid));
            }
            return ec;
        }
        break;

        default:
            break;
    }

    assert(0 && "unreachable here");
    return E_INVALID_ARGUMENT;
}

void CDebuggerStub::Dispose()
{
    delete this;
}

ECode CDebuggerStub::WriteProcessMemory(
    /* [in] */ HANDLE hProcess,
    /* [in] */ PVoid lpBaseAddress,
    /* [in] */ PCVOID lpBuffer,
    /* [in] */ size_t nSize,
    /* [out] */ size_t * lpNumberOfBytesWritten)
{
    DWORD dwOldProtect;
    BOOL ret = VirtualProtectEx(hProcess, lpBaseAddress, nSize,
                    PAGE_READWRITE, &dwOldProtect);
    if (!ret) {
        return GetLastErrorCode();
    }

    ret = ::WriteProcessMemory(
            hProcess, lpBaseAddress, lpBuffer, nSize, lpNumberOfBytesWritten);
    if (!ret) {
        return GetLastErrorCode();
    }

    ret = VirtualProtectEx(
            hProcess, lpBaseAddress, nSize, dwOldProtect, &dwOldProtect);
    if (!ret) {
        return GetLastErrorCode();
    }

    ret = FlushInstructionCache(hProcess, lpBaseAddress, nSize);
    if (!ret) {
        return GetLastErrorCode();
    }

    return NOERROR;
}

ECode CDebuggerStub::GetThreadStackBase(
    /* [in] */ DWORD dwProcessId,
    /* [in] */ DWORD dwThreadId,
    /* [in] */ PVoid * ppStackBase)
{
    ECode ec;

    if (NULL == ppStackBase) {
        return E_INVALID_ARGUMENT;
    }

    // Retrieve the fs segment register which gives a pointer to the TIB
    // for the target thread
    //
    HANDLE hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, dwThreadId);
    if (NULL == hThread) {
        return GetLastErrorCode();
    }

    Context ctx;
    ctx.ContextFlags = CONTEXT_ALL;
    LDT_ENTRY desc;

    BOOL ret = GetThreadContext(hThread, &ctx);
    if (!ret) {
        ec = GetLastErrorCode();
        CloseHandle(hThread);
        return ec;
    }

    ret = GetThreadSelectorEntry(hThread, ctx.SegFs, &desc);
    if (!ret) {
        ec = GetLastErrorCode();
    }
    else {
        ec = NOERROR;
    }

    CloseHandle(hThread);

    if (FAILED(ec)) {
        return ec;
    }

    // Read the thread information block from the memory of the target process
    //
    PVoid pAddress = (PVoid)((char *)(desc.BaseLow
                    | desc.HighWord.Bytes.BaseMid << 16
                    | desc.HighWord.Bytes.BaseHi << 24)
                    + offsetof(NT_TIB, StackBase));

    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwProcessId);
    if (NULL == hProcess) {
        return GetLastErrorCode();
    }

    ret = ReadProcessMemory(hProcess, pAddress, (void *)ppStackBase,
            sizeof(*ppStackBase), NULL);
    if (!ret) {
        ec = GetLastErrorCode();
    }
    else {
        ec = NOERROR;
    }

    CloseHandle(hProcess);

    return ec;
}

ECode CDebuggerStub::GetLastErrorCode()
{
    DWORD dwError = GetLastError();

    switch (dwError) {
        case ERROR_SUCCESS:         return NOERROR;
        case ERROR_ACCESS_DENIED:   return E_ACCESS_DENIED;
        case ERROR_INVALID_HANDLE:  return E_INVALID_ARGUMENT;
        case ERROR_SEM_TIMEOUT:     return E_TIMED_OUT;
        default:                    break;
    }
    return E_INVALID_OPERATION;
}

ECode CDebuggerStubManager::Read(
    /* [in] */ Int64 u64Offset,
    /* [in] */ Int32 bytesToRead,
    /* [out] */ MemoryBuf * pBuffer,
    /* [out] */ IEvent * * ppCompletionEvent)
{
    return E_NOT_IMPLEMENTED;
}

ECode CDebuggerStubManager::Write(
    /* [in] */ Int64 u64Offset,
    /* [in] */ const MemoryBuf & buffer,
    /* [out] */ Int32 * pBytesWritten,
    /* [out] */ IEvent * * ppCompletionEvent)
{
    return E_NOT_IMPLEMENTED;
}

ECode CDebuggerStubManager::Control(
        /* [in] */ Handle32 nControlCode,
        /* [in] */ const MemoryBuf & inBuffer,
        /* [out] */ MemoryBuf * pOutBuffer,
        /* [out] */ IEvent * * ppCompletionEvent)
{
    ECode ec = NOERROR;
    wchar_t * pServiceName;
    CDebuggerStub * pDbgStub = NULL;

    switch(nControlCode) {
        case DBGMANAGER_CODE_CREATE_DBGSTUB: {
            if (inBuffer.IsNullOrEmpty()) {
                return E_INVALID_ARGUMENT;
            }

            pServiceName = (wchar_t *)inBuffer.GetPayload();

            pDbgStub = new CDebuggerStub;
            if (NULL == pDbgStub) {
                return E_OUT_OF_MEMORY;
            }
            pDbgStub->AddRef();

            ec = _CSystem_RegisterRunningObject(pServiceName, pDbgStub);
            if (FAILED(ec)) {
                pDbgStub->Release();
            }
        }
        break;

        case DBGMANAGER_CODE_DESTROY_DBGSTUB: {
            if (inBuffer.IsNullOrEmpty()) {
                return E_INVALID_ARGUMENT;
            }

            pServiceName = (wchar_t *)inBuffer.GetPayload();

            ec = _CSystem_FindRunningObject(pServiceName, (IInterface**)&pDbgStub);
            if (FAILED(ec)) {
                return ec;
            }

            pDbgStub->Control(
                DBGSTUB_CTRL_CODE_DETACH_PROCESS, NULL_MEMORYBUF, NULL, NULL);

            _CSystem_UnregisterRunningObject(pServiceName);
            pDbgStub->Release();
        }
        break;

        default:
            ec = E_INVALID_ARGUMENT;
            break;
    }

    return ec;
}

void CDebuggerStubManager::Dispose()
{
}

IDeviceDriver * CDECL CreateDebuggerStubManager()
{
    static CDebuggerStubManager s_dbgStubManager;
    return &s_dbgStubManager;
}
